package com.zhxq.ai.repository;

import com.zhxq.common.config.RuoYiConfig;
import com.zhxq.common.core.domain.AjaxResult;
import com.zhxq.common.utils.file.FileUploadUtils;
import com.zhxq.common.utils.file.FileUtils;
import com.zhxq.common.utils.uuid.UUID;
import com.zhxq.framework.config.ServerConfig;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.pdfbox.Loader;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;
import org.springframework.web.multipart.MultipartFile;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.nio.file.Paths;
import java.util.List;
import java.util.Properties;

@Slf4j
@Component
@RequiredArgsConstructor
public class LocalPdfFileRepository  implements FileRepository{

    private final VectorStore vectorStore;

    private final ServerConfig serverConfig;

    // 会话id 与 文件名的对应关系，方便查询会话历史时重新加载文件
    private final Properties chatFiles = new Properties();

    @Override
    public AjaxResult save(String chatId, MultipartFile file) {
        try
        {
            // 文件类型验证
            if (!file.getOriginalFilename().endsWith(".pdf")) {
                return AjaxResult.error("仅支持PDF格式文件上传");
            }
            // 上传文件路径
            String filePath = RuoYiConfig.getUploadPath();
            // 上传并返回新文件名称
            String fileName = FileUploadUtils.upload(filePath, file);
            String url = serverConfig.getUrl() + fileName;
            String baseParentPath = Paths.get(filePath).getParent().toString();
            String cleanedRelativePath = fileName.replaceFirst("[\\\\/]?profile", "");
            // 获取本地pdf文件路径
            String fullPath = Paths.get(baseParentPath, cleanedRelativePath)
                    .normalize()
                    .toString();
            Resource resource = new FileSystemResource(fullPath);
            // 写入向量库
            writeToVectorStore(resource);
            // 存入会话id 与 文件名的对应关系
            chatFiles.put(chatId, resource.getFilename());
            File uploadedPdfFile = new File(fullPath);
            String coverImagePath = filePath + "/pdf_cover" + "/cover_" + UUID.randomUUID();
            // 生成封面图
            File coverImageFile = convertFirstPageToImage(uploadedPdfFile, coverImagePath);
            // 获取封面图URL
            String coverImageUrl = serverConfig.getUrl() + "/profile/upload/pdf_cover/" + coverImageFile.getName();
            int pages = getPages(uploadedPdfFile);
            AjaxResult ajax = AjaxResult.success();
            ajax.put("url", url);
            ajax.put("fileName", fileName);
            ajax.put("newFileName", FileUtils.getName(fileName));
            ajax.put("originalFilename", file.getOriginalFilename());
            ajax.put("coverImageUrl", coverImageUrl); // 返回封面图URL
            ajax.put("pages", pages);
            return ajax;
        }
        catch (Exception e)
        {
            return AjaxResult.error(e.getMessage());
        }
    }

    /**
     * 将PDF第一页转为图片
     *
     * @param pdfFile PDF文件路径
     * @param outputImagePath 输出图片路径（不带后缀）
     * @return 图片文件对象
     */
    public File convertFirstPageToImage(File pdfFile, String outputImagePath) throws IOException {
        try (PDDocument document = Loader.loadPDF(pdfFile)) {
            PDFRenderer renderer = new PDFRenderer(document);
            BufferedImage image = renderer.renderImage(0, 1.75f); // 第一页索引为 0，缩放比例 1.75x（相当于 ~96 DPI）
            File imageFile = new File(outputImagePath + ".png");
            ImageIO.write(image, "png", imageFile);
            return imageFile;
        } catch (IOException e) {
            throw new IOException("无法转换 PDF 首页为图片", e);
        }
    }

    /**
     * 获取PDF页数
     * @param pdfFile
     * @return
     * @throws IOException
     */
    public int getPages(File pdfFile) throws IOException {
        try (PDDocument document = Loader.loadPDF(pdfFile)) {
            return document.getNumberOfPages();
        } catch (IOException e) {
            throw new IOException("获取页数失败", e);
        }
    }

    @Override
    public Resource getFile(String chatId) {
        return new FileSystemResource(chatFiles.getProperty(chatId));
    }

    /**
     * 将pdf写入向量库
     * @param resource
     */
    private void writeToVectorStore(Resource resource) {
        // 1.创建PDF的读取器
        PagePdfDocumentReader reader = new PagePdfDocumentReader(
                resource, // 文件源
                PdfDocumentReaderConfig.builder()
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
                        .withPagesPerDocument(1) // 每1页PDF作为一个Document
                        .build()
        );
        // 2.读取PDF文档，拆分为Document
        List<Document> documents = reader.read();
        // 3.写入向量库
        vectorStore.add(documents);
    }
}
